import numpy as np

# import os
# import pandas as pd

input_file_path = 'EX_47108'
output_file_path = 'OUTPUT'

# Dataframe input file version
"""
os.chdir('C:/Users/owner/Desktop/Drought_NLP/Korea')
read = pd.read_csv('INPUT FILE')

yr = np.array(read['YR'], dtype='int')
mo = np.array(read['MO'], dtype='int')
dy = np.array(read['DY'], dtype='int')
prec = np.array(read['P'], dtype='float')
"""

acc_num = 365
Tday = 365

read = np.loadtxt(input_file_path, dtype='str')
yr = np.array(read[:, 0], dtype='int')
mo = np.array(read[:, 1], dtype='int')
dy = np.array(read[:, 2], dtype='int')
prec = np.array(read[:, 3], dtype='float')

s_point = yr[0] + 30

MEAN_syr = np.zeros((len(prec)))
MEAN_eyr = np.zeros((len(prec)))

for i in range(365, len(prec)):
    MEAN_syr[i] = yr[i] - 30
    MEAN_eyr[i] = yr[i] - 1
    if yr[i] < s_point:
        MEAN_syr[i] = yr[0]
        MEAN_eyr[i] = yr[0] + 29

EP = np.zeros((len(prec)))
MEP = np.zeros((len(prec)))
STD = np.zeros((len(prec)))
CDD = np.zeros((len(prec)))
DEP = np.zeros((len(prec)))
AvePrec = np.zeros((Tday))

SD = acc_num
imo = 0
BA = 0
CN = 0

for i in range(365, len(prec)):
    AP = 0
    VAR = 0
    for N in range(SD):
        AP = AP + prec[i - N]
        VAR = VAR + (AP / float(N + 1))
    EP[i] = VAR

imo = 0
for i in range(365, len(prec)):
    if imo == Tday:
        imo = 0
    C = MEAN_eyr[i] - MEAN_syr[i] + 1
    for l in range(Tday):
        A = 0.0
        for j in range(int(MEAN_syr[i]), int(MEAN_eyr[i]) + 1):
            k = (j - yr[0]) * Tday + l
            A = A + prec[int(k)]
        AvePrec[l] = A / C

    A = 0.0
    B = 0.0
    for j in range(Tday):
        m = imo - j
        if m < 0:
            m = m + Tday
        A += AvePrec[m]
        B += A / float(j + 1)
    MEP[i] = B

    A = 0.0
    for j in range(int(MEAN_syr[i]), int(MEAN_eyr[i]) + 1):
        k = (j - yr[0]) * Tday + imo
        A = A + (EP[k] - MEP[i]) ** 2
    C = MEAN_eyr[i] - MEAN_syr[i]
    STD[i] = np.sqrt(A / float(C))
    imo += 1

B = 0
C = 0
for i in range(365, len(prec)):
    A = EP[i] - MEP[i]
    if (A < 0) & (B < 0):
        C += 1
        CDD[i] = C
    else:
        C = 0
        CDD[i] = C
    B = A

print('EP, MEP, actual DS, and STD calculation complete')

scEDI = np.zeros((len(prec)))
imo = 0
aline = ''
aline = aline + 'YEAR MONTH DAY EDI Precipitation' + '\n'
start = 0
AvePrec = np.zeros((Tday))

for i in range(365, len(prec)):
    if imo == Tday:
        imo = 0
    if CDD[i] == 0:
        CDDEP = EP[i]
        CDDMEP = MEP[i]
        CDDSTD = STD[i]
    else:
        SD = CDD[i] + acc_num
        AP = 0
        VAR = 0
        for N in range(int(SD)):
            if (i - N) > 0:
                AP += prec[i - N]
                VAR += AP / float(N + 1)
        CDDEP = VAR
        C = MEAN_eyr[i] - MEAN_syr[i] + 1
        for m in range(Tday):
            A = 0.0
            for j in range(int(MEAN_syr[i]), int(MEAN_eyr[i]) + 1):
                k = (j - yr[0]) * Tday + m
                A += prec[k]
            AvePrec[m] = A / float(C)

        A = 0
        B = 0
        for j in range(int(SD)):
            m = imo - j
            if m < 0:
                for mx in range(999):
                    m = m + Tday
                    if m >= 0:
                        break
            A += AvePrec[m]
            B += A / float(j + 1)
        CDDMEP = B

        A = 0.0
        B = 0.0
        for j in range(int(MEAN_syr[i]), int(MEAN_eyr[i]) + 1):
            k = (j - yr[0]) * Tday + imo
            AP = 0
            VAR = 0
            for N in range(int(SD)):
                if (k - N) > 0:
                    AP += prec[k - N]
                    VAR += AP / float(N + 1)
            B = VAR
            A += (B - CDDMEP) ** 2
        C = MEAN_eyr[i] - MEAN_syr[i]
        CDDSTD = np.sqrt(A / float(C))
    imo += 1

    scEDI[i] = (CDDEP - CDDMEP) / float(CDDSTD)
    if (yr[i] == (yr[0] + 30)) & (dy[i] == dy[0]) & (mo[i] == mo[0]):
        start += 1
    if start > 0:
        aline = aline + str(yr[i]) + ' ' + str(mo[i]) + ' ' + str(dy[i]) + \
                ' ' + str(round(scEDI[i], 2)) + ' ' + str(prec[i]) + '\n'

oidx = open(output_file_path, 'w')
oidx.write(aline)
oidx.close()

print('scEDI calculation complete')